#!/usr/bin/env python
# -*- coding: utf-8 -*-
import torch
import torchvision
from torchvision import *
from torch import *
from PIL import Image
import numpy as np
from model import *
import matplotlib.pyplot as plt;

#打开目录，遍历文件
def open_dir(path):
    import os
    files = os.listdir(path)
    return files

#预测文件
def predict_file(path):
    #读取要预测的图片
    try:
        img = Image.open(path).convert('RGB')
        plt.imshow(img)
        plt.axis('off')
        plt.show()
    except:
        print('图片打开失败')


    trans = torchvision.transforms.Compose([
        torchvision.transforms.Resize((32, 32)),
        torchvision.transforms.ToTensor(),
        torchvision.transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))
    ])
    img = trans(img)
    img = img.to(device)
    img = img.unsqueeze(0)

    classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')
    output = model(img)
    prob = torch.nn.functional.softmax(output, dim=1)
    #print('概率: {}'.format(prob))
    value, predict = torch.max(output, 1)
    #print(predict.item())
    predict = output.argmax(dim=1)
    pred_class = classes[predict]
    print('预测结果: {}'.format(pred_class))

if __name__ == '__main__':
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = Modle()
    model.load_state_dict(torch.load('./model/model.pth'))
    model = model.to(device)
    #模式转为test模式
    model.eval()

    for i in open_dir('images'):
        print('预测图片: {}'.format(i))
        predict_file('images/'+i)
        print('------------------')


